#!/usr/bin/env python3
"""
Sample Target Generator
生成示例靶子图片用于测试目标检测功能
"""

import cv2
import numpy as np
import os
from pathlib import Path

def create_bullseye_target(size=640, rings=10):
    """创建一个标准的靶子图像"""
    # 创建白色背景
    img = np.ones((size, size, 3), dtype=np.uint8) * 255
    
    # 中心点
    center = (size // 2, size // 2)
    
    # 最大半径
    max_radius = size // 2 - 20
    
    # 绘制同心圆环
    colors = [(0, 0, 0), (255, 255, 255)]  # 黑白交替
    for i in range(rings, 0, -1):
        radius = int(max_radius * i / rings)
        color = colors[i % 2]
        cv2.circle(img, center, radius, color, -1)
    
    # 绘制红色靶心
    bullseye_radius = max_radius // 10
    cv2.circle(img, center, bullseye_radius, (0, 0, 255), -1)
    
    return img

def create_archery_target(size=640):
    """创建弓箭靶子"""
    img = np.ones((size, size, 3), dtype=np.uint8) * 255
    center = (size // 2, size // 2)
    max_radius = size // 2 - 20
    
    # 标准弓箭靶颜色 (从外到内)
    colors = [
        (255, 255, 255),  # 白色
        (0, 0, 0),        # 黑色
        (255, 255, 255),  # 白色
        (0, 0, 0),        # 黑色
        (0, 0, 255),      # 红色
        (255, 215, 0),    # 金色
        (0, 0, 255),      # 红色
        (255, 215, 0),    # 金色
        (0, 0, 255),      # 红色
        (255, 215, 0),    # 金色
    ]
    
    for i, color in enumerate(colors):
        radius = int(max_radius * (len(colors) - i) / len(colors))
        cv2.circle(img, center, radius, color, -1)
    
    return img

def create_training_labels(img_name, target_type="target"):
    """创建YOLO格式的标签文件"""
    # 假设整个图像就是一个靶子，边界框覆盖大部分图像
    # YOLO格式: class_id x_center y_center width height (归一化坐标)
    
    if target_type == "target":
        # 整个靶子
        label_content = "0 0.5 0.5 0.8 0.8\n"  # target class
        # 靶心
        label_content += "1 0.5 0.5 0.1 0.1\n"  # bullseye class
    
    return label_content

def generate_sample_dataset():
    """生成示例数据集"""
    # 创建目录
    base_path = Path("datasets/targets")
    train_img_path = base_path / "images" / "train"
    val_img_path = base_path / "images" / "val"
    train_label_path = base_path / "labels" / "train"
    val_label_path = base_path / "labels" / "val"
    
    # 生成训练集图片
    print("生成训练集图片...")
    for i in range(20):  # 生成20张训练图片
        # 生成不同类型的靶子
        if i % 2 == 0:
            img = create_bullseye_target(size=640, rings=10)
        else:
            img = create_archery_target(size=640)
        
        # 添加一些随机变化
        # 随机旋转
        angle = np.random.randint(-30, 30)
        M = cv2.getRotationMatrix2D((320, 320), angle, 1)
        img = cv2.warpAffine(img, M, (640, 640))
        
        # 随机亮度调整
        brightness = np.random.randint(-30, 30)
        img = cv2.convertScaleAbs(img, alpha=1, beta=brightness)
        
        # 保存图片
        img_name = f"target_{i:03d}.jpg"
        cv2.imwrite(str(train_img_path / img_name), img)
        
        # 生成标签
        label_content = create_training_labels(img_name)
        label_name = f"target_{i:03d}.txt"
        with open(train_label_path / label_name, 'w') as f:
            f.write(label_content)
    
    # 生成验证集图片
    print("生成验证集图片...")
    for i in range(5):  # 生成5张验证图片
        if i % 2 == 0:
            img = create_bullseye_target(size=640, rings=8)
        else:
            img = create_archery_target(size=640)
        
        # 保存图片
        img_name = f"val_target_{i:03d}.jpg"
        cv2.imwrite(str(val_img_path / img_name), img)
        
        # 生成标签
        label_content = create_training_labels(img_name)
        label_name = f"val_target_{i:03d}.txt"
        with open(val_label_path / label_name, 'w') as f:
            f.write(label_content)
    
    print(f"数据集生成完成！")
    print(f"训练集: 20张图片")
    print(f"验证集: 5张图片")
    print(f"数据集位置: {base_path}")

if __name__ == "__main__":
    generate_sample_dataset()